import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


img_dir = '../../../../../large_data/pic'
img_name = 'opencv-logo-white.png'
img_path = os.path.join(img_dir, img_name)

spr = 2
spc = 5
spn = 0
plt.figure(figsize=[12, 6])

img = cv.imread(img_path, cv.IMREAD_COLOR)
my_show_img(img, 'ori', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

SNR_list = [0.9, 0.7, 0.5, 0.3]
for SNR in SNR_list:
    img_noise = make_salt_and_pepper(img, SNR)
    my_show_img(img_noise, f'SNR = {SNR}', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
    img_blur = median_blur(img_noise, (5, 5))
    img_blur = np.uint8(img_blur)
    my_show_img(img_blur, f'Blured', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
